from Parameters import *
from PolymorphismTables import PolymorphismTables
from R2_Tools import R2
from SiteFrequencySpectrum import SiteFrequencySpectrum
import infiniteSitesTools as ist 
from Polyspect_Operator import Polyspect
from Polyspect_Power import PowerPoint


import math
import numpy.matlib as np
from scipy.stats import mstats
from scipy.stats import percentileofscore

import sys

def main(argv):
	
	print "***************************************************"
	print "********************* New run *********************"
	print "***************************************************"
	print ""
	

	pop = getDefaultPopulationParameters(50)
	
	popRatio = [1000.0, 100.0, 10.0, 1.0, 0.1, 0.01, 0.001, 0.0001]
	

	nullDemography = Demography( pop.N )
	nullDemography.addEpoch( pop.N, 0, True )
	
	simParams = SimulationParameters( pop, nullDemography )
	print simParams

	pt = PolymorphismTables()
	msCommand = simParams.getMSString_Null()
	msOut = pt.runMS(msCommand)
	pt.readMS(None, msOut)
	mySFSs = SiteFrequencySpectrum(pt)

	myTD = ist.getOp_TajimasD( pop.numChromosomes )
	TDpolyOp = Polyspect(myTD, pop, mySFSs )
	
	ssList = TDpolyOp.operateOnAll( TDpolyOp.mySFSs )
	print TDpolyOp.getRejectionValues(0.05)

	
	for pRatio in popRatio:
		tau = 1000
		alpha = 4.0 * pop.N * math.log(pRatio) / tau
		
		N1 = pop.N * math.exp(alpha * tau / (4.0*pop.N))
		
		altDemography_CC = Demography(pop.N)
		altDemography_CC.addEpoch( pop.N, 0, True)
		altDemography_CC.addEpoch( N1, tau, True )

		pt = PolymorphismTables()
		msCommand = simParams.getMSString_Alternative( altDemography_CC )
		msOut = pt.runMS(msCommand)
		pt.readMS(None, msOut)
		CC_SFSs = SiteFrequencySpectrum(pt)

		CC_Power = PowerPoint( TDpolyOp, CC_SFSs )
		print CC_Power.getPower()

		
		altDemography_CG = Demography(pop.N)
		altDemography_CG.addEpoch( alpha, 0, False)
		altDemography_CG.addEpoch( N1, tau, True )

		pt = PolymorphismTables()
		msCommand = simParams.getMSString_Alternative( altDemography_CG )
		msOut = pt.runMS(msCommand)
		pt.readMS(None, msOut)
		CG_SFSs = SiteFrequencySpectrum(pt)

		CG_Power = PowerPoint( TDpolyOp, CG_SFSs )

		print "N1\t" + str(N1) + "\tCC\t" + str(CC_Power.getPower()) + "\tCG\t" + str(CG_Power.getPower())



def getDefaultPopulationParameters(s=30):
	
	popParameters = Population()
	
	popParameters.numChromosomes = 28
	popParameters.numSegregatingSites = s

	popParameters.rateMutation = 1E-8
	popParameters.rateRecombination = 0.0
	popParameters.sequenceLength = 25000

	popParameters.N = 10000
	popParameters.numSimulations = 25000
	popParameters.numPowerSimulations = 10000

	return popParameters


if __name__ == "__main__":
	main(1) #sys.argv[1:])
